// (I)FFT(R)
// D. TAYLOR 2014

#ifndef __FFT_H__
#define __FFT_H__

#include "config.h"
#include "intrinsics.h"
#include "fpmath.h"

typedef struct {
  int32_t r;
  int32_t i;
} fft_complex_t;

// Set DC for symmetric (real) data
__attribute__((always_inline)) inline
void fft_sym_dc(fft_complex_t data[], int32_t r, int32_t i) {
  data[0].r = r + i;
  data[0].i = r - i;
}

// Process real data into complex symmetrical form
void fft_symreal(fft_complex_t data[], unsigned bits, bool bitreversed, bool invert) {
  unsigned size = 1 << --bits;
  unsigned shift = SINE_BITS - bits++;
  unsigned n, z, nc, zc;
  int32_t rsum, rdif, isum, idif;
  int32_t itwiddled, rtwiddled;
  for(nc = zc = size; nc; nc--, zc++) {
    if(bitreversed) {
      n = RBITS(nc, bits);
      z = RBITS(zc, bits);
    } else {
      n = nc; z = zc;
    }
    rsum = data[n].r + data[z].r;
    rdif = data[n].r - data[z].r;
    isum = data[n].i + data[z].i;
    idif = data[n].i - data[z].i;
    int32_t r =  sinetable[(1 << SINE_BITS) - (nc << shift)];
    int32_t i = -sinetable[nc << shift];
    if(invert) r = -r;
    rtwiddled = smmlar(r, isum, smmulr(i, rdif)) << 1;
    itwiddled = smmlsr(i, isum, smmulr(r, rdif)) << 1;
    data[n].r = rsum + rtwiddled;
    data[n].i = itwiddled + idif;
    data[z].r = rsum - rtwiddled;
    data[z].i = itwiddled - idif;
  }
  int32_t data_0_tr = data[0].r;
  data[0].r = (data[0].r + data[0].i);
  data[0].i = (data_0_tr - data[0].i);
  if(!invert) {
    data[0].r <<= 1;
    data[0].i <<= 1;
  }
}

void fft_bitreverse(fft_complex_t data[], unsigned bits) {
  unsigned size  = 1 << bits;
  unsigned shift = 32 - bits;
  for(uint32_t i = 1; i < size - 1; i++) {
    uint32_t z = rbit(i) >> shift;
    if(z > i) {
      int32_t t = data[i].r;
      data[i].r = data[z].r;
      data[z].r = t;
      t = data[i].i;
      data[i].i = data[z].i;
      data[z].i = t;
    }
  }
}

void fft_forward(fft_complex_t data[], unsigned size) {
  unsigned shift = SINE_BITS + 1;
  for(uint32_t step = 2 ; step <= size; step = step * 2, shift--) {
    uint32_t step2 = step >> 1;
    uint32_t step4 = step2 >> 1;
    for(uint32_t k = 0; k < step4 + (step2 & 1); k++) {
      int32_t twr = sinetable[(1 << SINE_BITS) - (k << shift)];
      int32_t twi = sinetable[k << shift];
      for(int32_t block = k; block < k + size; block += step) {
        uint32_t bstep = block + step2;
        int32_t r2 = data[bstep].r, i2 = data[bstep].i;
        int32_t r = smmlar(i2, twi, smmulr(r2,  twr));
        int32_t i = smmlar(i2, twr, smmulr(r2, -twi));
        int32_t r1 = data[block].r >> 1, i1 = data[block].i >> 1;
        data[block].r = r1 + r;
        data[block].i = i1 + i;
        data[bstep].r = r1 - r;
        data[bstep].i = i1 - i;
      }
    }
    for(uint32_t k = step2 & 1; k < step2 - step4; k++) {
      int32_t twr = -sinetable[k << shift];
      int32_t twi = sinetable[(1 << SINE_BITS) - (k << shift)];
      for(int32_t block = k + step4; block < k + step4 + size; block += step) {
        uint32_t bstep = block + step2;
        int32_t r2 = data[bstep].r, i2 = data[bstep].i;
        int32_t r = smmlar(i2, twi, smmulr(r2,  twr));
        int32_t i = smmlar(i2, twr, smmulr(r2, -twi));
        int32_t r1 = data[block].r >> 1, i1 = data[block].i >> 1;
        data[block].r = r1 + r;
        data[block].i = i1 + i;
        data[bstep].r = r1 - r;
        data[bstep].i = i1 - i;
      }
    }
  }
}

void fft_inverse(fft_complex_t data[], unsigned size) {
  unsigned shift = SINE_BITS + 1;
  for(unsigned step = 2; step <= size; step <<= 1, shift--) {
    uint32_t step2 = step >> 1;
    uint32_t step4 = step >> 2;
    for(uint32_t k = 0; k < step4 + (step2 & 1); k++) {
      int32_t twr = sinetable[(1 << SINE_BITS) - (k << shift)];
      int32_t twi = sinetable[k << shift];
      for(int32_t block = k; block < k + size; block += step) {
        uint32_t bstep = block + step2;
        //printf("%i %i\n", block, bstep);
        int32_t r1 = data[block].r, i1 = data[block].i;
        int32_t r2 = data[bstep].r, i2 = data[bstep].i;
        int32_t r = smmlar(i2, -twi, smmulr(r2, twr)) << 1;
        int32_t i = smmlar(i2, twr, smmulr(r2, twi)) << 1;
        data[block].r = r1 + r;
        data[block].i = i1 + i;
        data[bstep].r = r1 - r;
        data[bstep].i = i1 - i;
      }
    }
    for(uint32_t k = step2 & 1; k < step2 - step4; k++) {
      int32_t twr = -sinetable[k << shift];
      int32_t twi = sinetable[(1 << SINE_BITS) - (k << shift)];
      for(int32_t block = k + step4; block < k + step4 + size; block += step) {
        uint32_t bstep = block + step2;
        //printf("%i %i\n", block, bstep);
        int32_t r1 = data[block].r, i1 = data[block].i;
        int32_t r2 = data[bstep].r, i2 = data[bstep].i;
        int32_t r = smmlar(i2, -twi, smmulr(r2, twr)) << 1;
        int32_t i = smmlar(i2,  twr, smmulr(r2, twi)) << 1;
        data[block].r = r1 + r;
        data[block].i = i1 + i;
        data[bstep].r = r1 - r;
        data[bstep].i = i1 - i;
      }
    }
  }
}

// Magnitude and phase => complex vector for FFT
__inline
void fft_phase_magnitude(fft_complex_t complex[], unsigned bits, unsigned index, int32_t mag, uint32_t pha) {
  unsigned n = RBITS(index, bits);
  complex[n].r = smmulr(mag, sine(pha));
  complex[n].i = smmulr(mag, cosine(pha));
}

// Magnitude (phase==0) => complex vector for FFT
__inline
void fft_magnitude(fft_complex_t complex[], unsigned bits, unsigned index, int32_t mag) {
  unsigned n = RBITS(index, bits);
  complex[n].r = 0;
  complex[n].i = mag;
}

// REAL Symmetric magnitude and phase => complex vector for FFT
void fft_sym_phase_magnitude(fft_complex_t complex[], unsigned bits, unsigned index, int32_t mag_lo, int32_t pha_lo, int32_t mag_hi, int32_t pha_hi) {
  unsigned size = 1 << bits;
  unsigned shift = SINE_BITS - (bits - 1);
  unsigned n = RBITS(index, bits);
	unsigned z = RBITS(size - index, bits);
  int32_t rsum, rdif, isum, idif, r, i;
  int32_t itwiddled, rtwiddled;
  r = smmulr(mag_lo,   sine(pha_lo));
  i = smmulr(mag_hi,   sine(pha_hi));
  rsum = r + i; rdif = r - i;
  r = smmulr(mag_lo, cosine(pha_lo));
  i = smmulr(mag_hi, cosine(pha_hi));
  isum = r + i; idif = r - i;
  r = -sinetable[(1 << SINE_BITS) - (index << shift)];
  i = -sinetable[index << shift];
  rtwiddled = smmlar(r, isum, smmulr(i, rdif)) << 1;
  itwiddled = smmlsr(i, isum, smmulr(r, rdif)) << 1;
  complex[n].r = rsum + rtwiddled;
  complex[n].i = itwiddled + idif;
  complex[z].r = rsum - rtwiddled;
  complex[z].i = itwiddled - idif;
}

// REAL Symmetric magnitude (phase==0) => complex vector for FFT
void fft_sym_magnitude(fft_complex_t complex[], unsigned bits, unsigned index, int32_t mag_lo, int32_t mag_hi) {
  unsigned shift = SINE_BITS - bits + 1;
  unsigned n = RBITS(index, bits);
	unsigned z = RBITS((1 << bits) - index, bits);
  int32_t isum, idif, r, i;
  int32_t itwiddled, rtwiddled;
  isum = -(mag_lo + mag_hi); idif = mag_lo - mag_hi;
  r = sinetable[(1 << SINE_BITS) - (index << shift)];
  i = sinetable[index << shift];
  rtwiddled = smmulr(r, isum) << 1;
  itwiddled = smmulr(i, isum) << 1;
  complex[n].r = 0 + rtwiddled;
  complex[n].i = itwiddled + idif;
  complex[z].r = 0 - rtwiddled;
  complex[z].i = itwiddled - idif;
}

// Perform forward FFT
__inline
void fft_fft(fft_complex_t *complex, unsigned bits, bool bitreversed) {
  if(!bitreversed) fft_bitreverse(complex, bits);
  fft_forward(complex, 1 << bits);
}

// Perform inverse FFT
__inline
void fft_ifft(fft_complex_t *complex, unsigned bits, bool bitreversed) {
  if(!bitreversed) fft_bitreverse(complex, bits);
  fft_inverse(complex, 1 << bits);
}

// Perform forward FFT on REAL data
__inline
void fft_fftr(fft_complex_t *complex, unsigned bits, bool bitreversed) {
  fft_fft(complex, bits, bitreversed);
  fft_symreal(complex, bits, false, false);
}

// Perform inverse FFT on REAL data
__inline
void fft_ifftr(fft_complex_t *complex, unsigned bits, bool bitreversed) {
  fft_symreal(complex, bits, bitreversed, true);
  fft_ifft(complex, bits, bitreversed);
}

#endif
